#!/usr/bin/python
"""
This program controls the chassis fan speed through PWM based on the temperature
of the hottest hard drive in the chassis. It uses the IBM M1015 or LSI tool
'MegaCli' for reading hard drive temperatures.
"""
import os
import sys
import subprocess
import re
import time
import syslog
import multiprocessing as mp
import copy_reg
import types
import ConfigParser

def _reduce_method(meth):
    """
    This is a hack to work around the fact that multiprocessing
    can't operate on class methods by default.
    """
    return (getattr, (meth.__self__, meth.__func__.__name__))

class PID:
    """
    Discrete PID control
    Source: http://code.activestate.com/recipes/577231-discrete-pid-controller/

    This class calculates the appropriate fan speed based on the difference
    between the current temperature and the desired (target) temperature.
    """

    def __init__(self, P, I, D, Derivator, Integrator, \
                Integrator_max, Integrator_min):
        """
        Generic initialisation of local variables.
        """
        self.Kp = P
        self.Ki = I
        self.Kd = D
        self.Derivator = Derivator
        self.Integrator = Integrator
        self.Integrator_max = Integrator_max
        self.Integrator_min = Integrator_min

        self.set_point = 0.0
        self.error = 0.0

    def update(self, current_value):
        """
        Calculate PID output value for given reference input and feedback
        Current_value = set_point - measured value (difference)
        """
        self.error = current_value - int(self.set_point)

        self.P_value = self.Kp * self.error
        self.D_value = self.Kd * ( self.error + self.Derivator)
        self.Derivator = self.error

        self.Integrator = self.Integrator + self.error

        if self.Integrator > self.Integrator_max:
            self.Integrator = self.Integrator_max
        elif self.Integrator < self.Integrator_min:
            self.Integrator = self.Integrator_min

        self.I_value = self.Integrator * self.Ki

        PID = self.P_value + self.I_value + self.D_value

        return PID

    def set_target_value(self, set_point):
        """
        Initilize the setpoint of PID
        """
        self.set_point = set_point

copy_reg.pickle(types.MethodType, _reduce_method)
class Smart:
    """
    Uses SMART data from storage devices to determine the temperature
    of the hottest drive.
    """

    def __init__(self):
        """
        Init.
        """
        self.block_devices = ""
        self.device_filter = "sd"
        self.highest_temperature = 0
        self.get_block_devices()
        self.smart_workers = 24

    def filter_block_devices(self, block_devices):
        """
        Filter out devices like 'loop, ram'.
        """
        devices = []
        for device in block_devices:
            if not device.find(self.device_filter):
                devices.append(device)
        return devices

    def get_block_devices(self):
        """
        Retrieve the list of block devices.
        By default only lists /dev/sd* devices.
        Configure the appropriate device filter with
        setting <object>.device_filter to some other value.
        """
        devicepath = "/sys/block"
        block_devices = os.listdir(devicepath)
        block_devices.sort()
        self.block_devices = self.filter_block_devices(block_devices)

    def get_smart_data(self, device):
        """
        Call the smartctl command line utilily on a device to get the raw
        smart data output.
        """

        device = "/dev/" + device

        try:
            child = subprocess.Popen(['smartctl',  '-a',  '-d',  'ata', \
                                    device], stdout=subprocess.PIPE, \
                                     stderr=subprocess.PIPE)
        except OSError:
            print "Executing smartctl gave an error,"
            print "is smartmontools installed?"
            sys.exit(1)

        rawdata = child.communicate()

        if child.returncode:
            child = subprocess.Popen(['smartctl',  '-a',  device],
                                     stdout=subprocess.PIPE,
                                     stderr=subprocess.PIPE)
            rawdata = child.communicate()
            if child.returncode == 1:
                return ""

        smartdata = rawdata[0]
        return smartdata

    def get_parameter_from_smart(self, data, parameter, distance):
        """
        Retreives the desired value from the raw smart data.
        """
        regex = re.compile(parameter + '(.*)')
        match = regex.search(data)

        if match:
            tmp = match.group(1)
            length = len(tmp.split("   "))
            if length <= distance:
                distance = length-1

            #
            # SMART data is often a bit of a mess,  so this
            # hack is used to cope with this.
            #

            try:
                model = match.group(1).split("   ")[distance].split(" ")[1]
            except:
                model = match.group(1).split("   ")[distance+1].split(" ")[1]
            return str(model)
        return 0

    def get_temperature(self, device):
        """
        Get the current temperature of a block device.
        """
        smart_data = self.get_smart_data(device)
        temperature = int(self.get_parameter_from_smart(smart_data, \
                                            'Temperature_Celsius', 10))
        return temperature

    def get_highest_temperature(self):
        """
        Get the highest temperature of all the block devices in the system.
        Because retrieving SMART data is slow, multiprocessing is used
        to collect SMART data in parallel from multiple devices.
        """
        highest_temperature = 0
        pool = mp.Pool(processes=int(self.smart_workers))
        results = pool.map(self.get_temperature, self.block_devices)
        pool.close()

        for temperature in results:
            if temperature > highest_temperature:
                highest_temperature = temperature
        self.highest_temperature = highest_temperature

        return self.highest_temperature

class Controller:
    """
    Reading temperature data from IBM / LSI controllers.
    """
    def __init__(self):
        self.megacli = "/opt/MegaRAID/MegaCli/megacli"
        self.ports_per_controller = 8
        self.highest_temperature = 0

    def number_of_controllers(self):
        """
        Get the number of LSI HBAs on the system.
        In my case, I have 3 controllers with 8 drives each.
        """
        rawdata = subprocess.Popen(\
            [self.megacli,'-cfgdsply','-aALL'],\
             stdout=subprocess.PIPE, stderr=subprocess.PIPE).communicate()[0]
        regex = re.compile('Adapter:.*')
        match = regex.findall(rawdata)
        return len(match)

    def get_drive_temp(self, controller, port):
        """
        Get the temperature from an individual drive through the megacli
        utility. The return value is a positive integer that specifies the
        temperature in Celcius.
        """
        rawdata =  subprocess.Popen(\
            [self.megacli,  '-pdinfo', '-physdrv', '[64:' +\
               str(port) +']', '-a' + str(controller)],\
               stdout=subprocess.PIPE, stderr=subprocess.PIPE).communicate()[0]

        regex = re.compile('Drive Temperature :(.*)')
        match = regex.search(rawdata)
        try:
            temp = match.group(1).split("C")[0]

            # Ugly hack: issue with some old WD drives
            # Controller reports 65C for them.
            if temp == "N/A":
                temp = "?"
            if int(temp) >= 60:
                temp = "?"
            return temp

        except(AttributeError):
            return ""
        except(IndexError):
            return ""

    def fetch_data(self):
        """
        Returns a two-dimentional list containing
        the temperature of each drive. The first dimension is the
        chassis. The second dimension is the drive.
        """
        drivearray = \
             [[0 for x in xrange(self.ports_per_controller)]\
                for x in xrange(self.number_of_controllers())]

        for controller in xrange(self.number_of_controllers()):
            for port in xrange(self.ports_per_controller):
                disk = self.get_drive_temp(controller, port)
                if len(disk) == 0:
                    disk = ""
                drivearray[controller][port] = disk

        return drivearray

    def get_highest_temperature(self):
        """
        Walks through the list of all the drives and compares
        all drive temperatures. The highest drive temperature
        is returned as an integer, representing degrees of Celcius.
        """

        data = self.fetch_data()
        temperature = 0
        for controller in data:
            for disk in controller:
                if disk > temperature:
                    temperature = disk
        self.highest_temperature = int(temperature)

        return self.highest_temperature


class FanControl:
    """
    The chassis object provides you with the option:
    1. Get the temperature of the hottest hard drive
    2. Get the current fan speed
    3. Set the fan speed
    """

    def __init__(self):
        """
        Generic init method.
        """
        self.polling_interval = 30
        self.pwm_max = 255
        self.pwm_min = 100
        self.pwm_safety = 160
        self.fan_speed = 50
        self.fan_control_enable = ""
        self.fan_control_device = ""
        self.debug = False

    def get_pwm(self):
        """
        Return the current PWM speed setting.
        """
        PWM=""

        for device in self.fan_control_device:
            filename = device
            filehandle = open(filename, 'r')
            pwm_value = int(filehandle.read().strip())
            filehandle.close()
            PWM = PWM + " " + str(pwm_value)

        return PWM

    def set_pwm(self, value):
        """
        Sets the fan speed. Only allows values between
        pwm_min and pwm_max. Values outside these ranges
        are set to either pwm_min or pwm_max as a safety
        precaution.
        """
        self.enable_fan_control()

        for device in self.fan_control_device:

            filename = device
            pwm_max = self.pwm_max
            pwm_min = self.pwm_min

            value = pwm_max if value > pwm_max else value
            value = pwm_min if value < pwm_min else value

            filehandle = open(filename, 'w')
            filehandle.write(str(value))
            filehandle.close()

    def set_fan_speed(self, percent):
        """
        Set fan speed based on a percentage of full speed.
        Values are thus 1-100 instead of raw 1-255
        """
        self.fan_speed = percent
        one_percent = float(self.pwm_max) / 100
        pwm = percent * one_percent
        self.set_pwm(int(pwm))

    def enable_fan_control(self):
        """
        Tries to enable manual fan speed control."
        """
        for device in self.fan_control_device:
            filename = device
            filehandle = open(filename, 'w')
            try:
                filehandle.write('1')
                filehandle.close()
            except IOError:
                message = "Error enabling fan control. Sufficient privileges?"
                print message
                sys.exit(1)


def is_debug_enabled():
    """
    Set debug if enabled.
    """
    try:
        debug = os.environ['DEBUG']
        if debug == "True":
            return True
        else:
            return False

    except (KeyError):
        return False

def log(temperature, chassis, pid):
    """
    Logging to syslog and terminal (export DEBUG=True).
    """
    P = str(pid.P_value)
    I = str(pid.I_value)
    D = str(pid.D_value)
    E = str(pid.error)

    TMP = str(temperature)
    PWM = str(chassis.get_pwm())
    PCT = str(chassis.fan_speed)

    all_vars = [TMP, PCT, PWM, P, I, D, E]
    formatstring = "Temp: {:2} | FAN: {:2}% | PWM: {:3} | P={:3} | I={:3} | "\
                   "D={:3} | Err={:3}|"

    msg = formatstring.format(*all_vars)

    syslog.openlog("SFC")
    syslog.syslog(msg)

    if is_debug_enabled():
        print msg

def read_config():
    """ Main"""
    config_file = "/etc/storagefancontrol"
    conf = ConfigParser.SafeConfigParser()
    conf.read(config_file)
    return conf

def get_pid_settings(config):
    """ Get PID settings """
    P = config.getint("Pid", "P")
    I = config.getint("Pid", "I")
    D = config.getint("Pid", "D")
    D_amplification = config.getint("Pid", "D_amplification")
    I_start = config.getint("Pid", "I_start")
    I_max = config.getint("Pid", "I_max")
    I_min = config.getint("Pid", "I_min")

    pid = PID(P, I, D, D_amplification, I_start, I_max, I_min)
    target_temperature = config.getint("General", "target_temperature")
    pid.set_target_value(target_temperature)

    return pid

def get_temp_source(config):
    """ Configure temperature source."""

    mode = config.get("General", "mode")

    if mode == "smart":
        temp_source = Smart()
        temp_source.device_filter = config.get("Smart", "device_filter")
        temp_source.smart_workers = config.getint("Smart", "smart_workers")
        return temp_source

    if mode == "controller":
        temp_source = Controller()
        temp_source.megacli = config.get("Controller", "megacli")
        temp_source.ports_per_controller = config.getint("Controller", \
                                         "ports_per_controller")
        return temp_source

    print "Mode not set, check config."
    sys.exit(1)

def get_chassis_settings(config):
    """ Initialise chassis fan settings. """

    chassis = FanControl()
    chassis.pwm_min = config.getint("Chassis", "pwm_min")
    chassis.pwm_max = config.getint("Chassis", "pwm_max")
    chassis.pwm_safety = config.getint("Chassis", "pwm_safety")
    chassis.fan_control_enable = config.get( "Chassis", "fan_control_enable")
    chassis.fan_control_enable = chassis.fan_control_enable.split(",")
    chassis.fan_control_device = config.get("Chassis", "fan_control_device")
    chassis.fan_control_device = chassis.fan_control_device.split(",")
    return chassis

def main():
    """
    Main function. Contains variables that can be tweaked to your needs.
    Please look at the class object to see which attributes you can set.
    The pid values are tuned to my particular system and may require
    ajustment for your system(s).
    """
    config = read_config()
    polling_interval = config.getfloat("General", "polling_interval")

    chassis = get_chassis_settings(config)
    pid = get_pid_settings(config)
    temp_source = get_temp_source(config)

    try:
        while True:
            highest_temperature = temp_source.get_highest_temperature()
            fan_speed = pid.update(highest_temperature)
            chassis.set_fan_speed(fan_speed)
            log(highest_temperature, chassis, pid)
            time.sleep(polling_interval)

    except (KeyboardInterrupt, SystemExit):
        chassis.set_pwm(chassis.pwm_safety)
        sys.exit(1)

if __name__ == "__main__":
    main()

